import torch
from torch import nn

class MemoryModule(nn.Module):
    def __init__(self):
        super().__init__()
    
    def reset(self):
        raise NotImplementedError()
    
class BaseNet(nn.Module):
    def __init__(self):
        super().__init__()

    def reset(self):
        
        def check_for_reset(module: nn.Module):
            for name, child in module.named_children():
                if isinstance(child, MemoryModule):
                    child.reset()
                elif list(child.children()):
                    check_for_reset(child)
        
        check_for_reset(self)